from termcolor import colored
import boto3
import sys
from datetime import datetime

author = {
    "name":"gl4ssesbo1",
    "twitter":"https://twitter.com/gl4ssesbo1",
    "github":"https://github.com/gl4ssesbo1",
    "blog":"https://www.pepperclipp.com/"
}

needs_creds = True

variables = {
    "SERVICE": {
        "value":"sts",
        "required":"true",
        "description":"The service that will be used to run the module. It cannot be changed."
    },
    "ROLEARN": {
        "value":"",
        "required":"true",
        "description":"The service that will be used to run the module. It cannot be changed."
    },
    "ROLESESSIONNAME": {
        "value":"",
        "required":"true",
        "description":"The service that will be used to run the module. It cannot be changed."
    },
    "POLICY": {
        "value":"",
        "required":"false",
        "description":"The service that will be used to run the module. It cannot be changed."
    },
    "EXTERNALID": {
        "value":"",
        "required":"false",
        "description":"The service that will be used to run the module. It cannot be changed."
    },
    "POLICYARNS": {
        "value":"",
        "required":"false",
        "description":"The service that will be used to run the module. It cannot be changed."
    },
    "DURATIONSECONDS": {
        "value":"900",
        "required":"true",
        "description":"The service that will be used to run the module. It cannot be changed."
    },
    "TOKENCODE": {
        "value":"",
        "required":"false",
        "description":"The service that will be used to run the module. It cannot be changed."
    }
}

description = "Get the ID of the current User. Just provide IAM Credentials and run. No extra permissions are needed"

aws_command = "aws sts assume-role --role-arn <role arn> --role-session-name <rolename> --profile <profile>"

def exploit(profile, workspace):
    try:
        now = datetime.now()
        dt_string = now.strftime("%d_%m_%Y_%H_%M_%S")
        file = "{}_sts_assume_role".format(dt_string)
        filename = "./workspaces/{}/{}".format(workspace, file)

        rolesessionname = variables['ROLESESSIONNAME']['value']
        rolearn = variables['ROLEARN']['value']
        durationseconds = int(variables['DURATIONSECONDS']['value'])

        policy = variables['POLICY']['value']
        tokencode = variables['TOKENCODE']['value']
        externalid = variables['EXTERNALID']['value']

        arguments = {}

        arguments['RoleArn'] = rolearn
        arguments['RoleSessionName'] = rolesessionname
        arguments['DurationSeconds'] = durationseconds

        if not tokencode == "":
            arguments['TokenCode'] = tokencode

        if not policy == "":
            arguments['Policy'] = policy

        if not externalid == "":
            arguments['ExternalId'] = externalid

        if not variables['POLICYARNS']['value'] == "":
            policyarns = []
            for parn in (variables['POLICYARNS']['value']).split(","):
                arn = {}
                arn['arn'] = parn
                policyarns.append(arn)

            arguments['PolicyArns'] = policyarns

        response = profile.assume_role(**arguments)
        
        full_response = {}
        for key,value in response.items():
            if not key == "ResponseMetadata":
                full_response[key] = value
        
        output = ""
        output += (colored("------------------------------------------------\n", "yellow", attrs=['bold']))
        output += ("{}: {}\n".format(colored("AccessKeyId", "yellow", attrs=['bold']), full_response['Credentials']['AccessKeyId']))
        output += (colored("------------------------------------------------\n", "yellow", attrs=['bold']))
        output += "\n"

        for key,value in full_response.items():
            if key == "PackedPolicySize":
                output += ("\t{}: {}\n".format(colored(key, "yellow", attrs=['bold']), colored(str(value), "red")))
            else:
                output += ("\t{}: \n".format(colored(key, "yellow", attrs=['bold'])))
                for k,v in value.items():
                    output += ("\t\t{}: {}\n".format(colored(k, "red", attrs=['bold']), colored(str(v), "blue")))
            
            output += "\n"
        print(output)
        
    except profile.exceptions.MalformedPolicyDocumentException as e:
        print(colored("[*] {}".format(e), "red"))

    except profile.exceptions.PackedPolicyTooLargeException as e:
        print(colored("[*] {}".format(e), "red"))

    except profile.exceptions.RegionDisabledException as e:
        print(colored("[*] {}".format(e), "red"))

    except profile.exceptions.ExpiredTokenException as e:
        print(colored("[*] {}".format(e), "red"))

    except:
        e = sys.exc_info()
        print(colored("[*] {}".format(e), "red"))